from PySide6.QtWidgets import QApplication, QMainWindow, QFileDialog, QMenu
from PySide6.QtGui import QImage, QPixmap, QColor
from PySide6.QtCore import QTimer, QThread, Signal, QObject, QPoint, Qt
from ui.CustomMessageBox import MessageBox
from ui.home import Ui_MainWindow
from UIFunctions import *
import numpy as np
from PIL import Image

import re
import cv2
import sys
import os
from deeplab import DeeplabV3

MODE_IMAGE = "image"
MODE_VIDEO = "video"
mode = MODE_IMAGE  # 默认模式为图片模式

class Predictor(QObject):
    # 定义信号
    pre_img_signal = Signal(np.ndarray)  # 原始图像信号
    res_img_signal = Signal(np.ndarray)  # 推理结果图像信号
    status_msg_signal = Signal(str)  # 状态信息信号
    fps_signal = Signal(str)  # 帧率信号
    progress_signal = Signal(int)  # 进度信号
    image_process_signal = Signal(np.ndarray)  # 新增信号，用于处理图像


    def __init__(self):
        super(Predictor, self).__init__()
        self.source = ''  # 输入图像路径
        self.stop = False  # 停止信号
        self.pause = True  # 暂停信号
        #self.model_path = "./best_epoch_weights.pth"  # 模型路径
        self.model_path = ""  # 模型路径
        #self.deeplab = DeeplabV3(model_path=self.model_path)  # 实例化模型
        self.deeplab = None  # 实例化模型
        self.name_classes = ["_background_", "leaf", "point"]  # 类别名称列表
        self.save_res = False

    def process_frame(self, img):
        # 处理图像的槽函数
        if not self.pause and not self.stop:
            try:
                # 运行推理
                print("ok")
                img = Image.fromarray(img)
                r_image = self.deeplab.detect_image(img, count=False, name_classes=self.name_classes)

                # 将推理结果转换为 BGR 格式
                r_image_rgb = cv2.cvtColor(np.array(r_image), cv2.COLOR_RGB2BGR)
                # 发射推理结果图像信号
                self.res_img_signal.emit(r_image_rgb)
                
                # 发送状态信息
                self.status_msg_signal.emit('Frame processed')
            except Exception as e:
                # 发送错误信息
                self.status_msg_signal.emit(f'Error: {e}')
                print(f'Error: {e}')


    def run_model(self):
        # 单张图片模式下的推理逻辑
        if mode == MODE_IMAGE:
            img_path = self.source
            try:
                # 使用 Pillow 加载图像
                image = Image.open(img_path)
                if image is None:
                    raise FileNotFoundError(f"Image not found: {img_path}")

                # 转换为 RGB 格式
                image_rgb = image.convert("RGB")
                # 发射原始图像信号
                self.pre_img_signal.emit(np.array(image_rgb))

                # 运行推理
                r_image = self.deeplab.detect_image(image_rgb, count=False, name_classes=self.name_classes)
                
                # 将推理结果转换为 BGR 格式
                r_image_rgb = cv2.cvtColor(np.array(r_image), cv2.COLOR_RGB2BGR)
                # 发射推理结果图像信号
                self.res_img_signal.emit(r_image_rgb)
                

                # 如果 save_res 为 True，则保存结果
                if self.save_res:
                    result_path = "./result.jpg"  # 你可以根据需要修改保存路径
                    cv2.imwrite(result_path, r_image_rgb)
                    print(f"Result image saved to: {result_path}")
                # 发送状态信息
                self.status_msg_signal.emit('Detection completed')
            except Exception as e:
                # 发送错误信息
                self.status_msg_signal.emit(f'Error: {e}')
                print(f'Error: {e}')


class MainWindow(QMainWindow, Ui_MainWindow):

    def __init__(self, parent=None):
        super(MainWindow, self).__init__(parent)
        
        # 设置基本UI
        self.setupUi(self)
        self.setAttribute(Qt.WA_TranslucentBackground)
        self.setWindowFlags(Qt.FramelessWindowHint)
        UIFuncitons.uiDefinitions(self)

        # 设置阴影
        UIFuncitons.shadow_style(self, self.Class_QF, QColor(162,129,247))

        UIFuncitons.shadow_style(self, self.Model_QF, QColor(64, 186, 193))
        

        # 初始化Timer用于监控模型文件
        self.Qtimer_ModelBox = QTimer(self)

        self.Qtimer_ModelBox.start(2000)

        # 预测器线程
        self.predictor = Predictor()
        self.thread = QThread()  # 创建线程
        self.predictor.moveToThread(self.thread)

        # 连接信号槽
        self.predictor.pre_img_signal.connect(lambda x: self.show_image(x, self.pre_video))
        self.predictor.res_img_signal.connect(lambda x: self.show_image(x, self.res_video))
        self.predictor.status_msg_signal.connect(self.show_status)
        self.predictor.fps_signal.connect(lambda x: self.fps_label.setText(x))
        self.predictor.progress_signal.connect(lambda x: self.progress_bar.setValue(x))
        
        self.predictor.image_process_signal.connect(self.predictor.process_frame)

        # UI按钮和滑动条等控件连接
        self.run_button.clicked.connect(self.run_or_continue)
        self.stop_button.clicked.connect(self.stop)
        self.save_res_button.toggled.connect(self.is_save_res)

        self.ToggleBotton.clicked.connect(lambda: UIFuncitons.toggleMenu(self, True))
        self.settings_button.clicked.connect(lambda: UIFuncitons.settingBox(self, True))
        self.src_file_button.clicked.connect(self.open_src_file)  # 点击按钮打开文件选择
        self.src_cam_button.clicked.connect(self.start_camera)     # 启动摄像头
        self.src_rtsp_button.clicked.connect(self.open_Model)     # 导入模型文件

        self.capture = None
        self.load_config()

    def open_Model(self):
        # 打开文件对话框选择以.pth结尾的文件
        
        name, _ = QFileDialog.getOpenFileName(self, 'Open File', "", "Model Files (*.pth)")
        
        if name:
            self.show_status(f'Load Model: {os.path.basename(name)}')
            # 假设load_model是加载.pth文件的函数，你可以根据需要实现此功能  
            self.predictor.model_path = name
            print(name)
            match = re.search(r'([^/\\]+)(?=\.pth$)', name).group(1)[:12]
            print(match)
            self.Model_name.setText(match)
            self.predictor.deeplab = DeeplabV3(model_path=name)


    def start_camera(self):
        # 启动摄像头
        if self.capture is None or not self.capture.isOpened():
            self.capture = cv2.VideoCapture(0)  # 重新初始化摄像头
        if not self.capture.isOpened():
            self.show_status("Failed to open camera")
            return
        self.timer = QTimer(self)  # 重新初始化定时器
        self.timer.timeout.connect(self.update_frame)  # 定时器超时时调用 update_frame 方法
        self.timer.start(30)  # 设置为30毫秒定时捕获
        global mode
        mode = MODE_VIDEO  # 设置模式为视频模式


    def update_frame(self):
        # 从摄像头捕获帧并显示
        ret, frame = self.capture.read()
        if ret:
            # 翻转图像
            frame = cv2.flip(frame, 1)  # 水平翻转图像
            # 转换为RGB格式
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            self.show_image(frame, self.pre_video)

            # 如果检测处于运行状态，进行推理
            if self.run_button.isChecked() and not self.predictor.pause:
                # 发射原始图像信号
                self.predictor.image_process_signal.emit(frame)


    def open_src_file(self):
        # 打开文件对话框选择图像或视频文件
        name, _ = QFileDialog.getOpenFileName(self, 'Open File', "", "Image Files (*.jpg *.png *.bmp);;Video Files (*.mp4 *.mkv *.avi)")
        
        if name:
            self.show_status(f'Load File: {os.path.basename(name)}')
            self.load_image(name)  # 加载选中的图像
            self.predictor.source = name  # 设置预测器的输入源


    def load_image(self, image_path):
        # 控制台输出路径
        print(f'Loading image from {image_path}')

        # 读取图像，这里读入的是BGR格式
        img = cv2.imread(image_path)
        if img is not None:
            # 将图像从BGR转换为RGB
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  
            self.show_image(img, self.pre_video)  # 在pre_video标签中显示图像
        else:
            self.show_status('Failed to load image.')

    @staticmethod
    def show_image(img_src, label):
        try:
            ih, iw, _ = img_src.shape
            w = label.geometry().width()
            h = label.geometry().height()
            # 保持原始数据比例
            if iw/w > ih/h:
                scale = w / iw
                nw = w
                nh = int(scale * ih)
                img_src_ = cv2.resize(img_src, (nw, nh))
            else:
                scale = h / ih
                nw = int(scale * iw)
                nh = h
                img_src_ = cv2.resize(img_src, (nw, nh))

            # 创建RGB格式的QImage
            img = QImage(img_src_.data, img_src_.shape[1], img_src_.shape[0], img_src_.shape[2] * img_src_.shape[1],
                        QImage.Format_RGB888)
            label.setPixmap(QPixmap.fromImage(img))  # 更新标签显示图像

        except Exception as e:
            print(repr(e))

    def show_status(self, msg):
        self.status_bar.setText(msg)  # 设置状态栏文本
        if msg == 'Detection completed' or msg == '检测完成':
            self.save_res_button.setEnabled(True)  # 允许保存结果图像

            self.run_button.setChecked(False)  # 停止运行按钮的勾选状态
            self.progress_bar.setValue(0)  # 重置进度条
            if self.thread.isRunning():  # 如果线程正在运行，则退出
                self.thread.quit()
        elif msg == 'Detection terminated!' or msg == '检测终止':
            self.save_res_button.setEnabled(True)  # 允许保存结果图像

            self.run_button.setChecked(False)  # 停止运行按钮的勾选状态
            self.progress_bar.setValue(0)  # 重置进度条
            if self.thread.isRunning():  # 如果线程正在运行，则退出
                self.thread.quit()
            self.pre_video.clear()  # 清除预览视频显示
            self.res_video.clear()  # 清除结果视频显示
            self.Class_num.setText('--')  # 重置类别数量显示


    def run_or_continue(self):
        if mode == MODE_IMAGE:
            if self.predictor.source == '':
                self.show_status('Please select a video source before starting detection...')
                self.run_button.setChecked(False)
            else:
                self.predictor.stop = False  # 重置停止标志
                if self.run_button.isChecked():  # 如果运行按钮被勾选
                    self.run_button.setChecked(True)  # 保持运行按钮的勾选状态

                    self.save_res_button.setEnabled(False)  # 禁用保存结果图像按钮
                    self.show_status('Detecting...')  # 显示检测中状态
                    self.predictor.pause = False  # 设置预测器为非暂停状态
                    if not self.thread.isRunning():  # 如果线程未运行
                        print("thread start")
                        self.thread.start()  # 启动线程
                        self.predictor.run_model()  # 调用预测器的run_model方法
                        print("thread end")
                else:
                    self.predictor.pause = True  # 设置预测器为暂停状态
                    self.show_status("Pause...")  # 显示暂停状态
                    self.run_button.setChecked(False)  # 取消运行按钮的勾选状态
        elif mode == MODE_VIDEO:
            if self.run_button.isChecked():
                self.run_button.setChecked(True)

                self.save_res_button.setEnabled(False)
                self.show_status('Detecting from camera...')
                self.predictor.pause = False
                if not self.thread.isRunning():
                    self.thread.start()
                    # 调用run_model进行视频流推理
                    self.predictor.run_model()
            else:
                self.predictor.pause = True
                self.show_status("Pause...")
                self.run_button.setChecked(False)


    def is_save_res(self):
        if self.save_res_button.checkState() == Qt.CheckState.Unchecked:
            self.show_status('NOTE: Run image results are not saved.')
            self.predictor.save_res = False  # 如果需要保存结果，可以在此保留逻辑
        elif self.save_res_button.checkState() == Qt.CheckState.Checked:
            self.show_status('NOTE: Run image results will be saved.')
            self.predictor.save_res = True  # 如果需要保存结果，可以在此保留逻辑
    


    def load_config(self):
        self.Class_num.setText('2')
        self.show_status("Welcome to the DeepLab integration!")

    def stop(self):
        # 停止线程
        if self.thread.isRunning():
            self.thread.quit()
        
        # 停止摄像头预览
        if hasattr(self, 'timer') and self.timer.isActive():
            self.timer.stop()  # 停止定时器
        if hasattr(self, 'capture') and self.capture.isOpened():
            self.capture.release()  # 释放摄像头资源
            self.capture = None  # 重置摄像头对象为 None

        # 重置模式为图片模式
        global mode
        mode = MODE_IMAGE

        # 重置预测器状态
        self.predictor.stop = True
        self.predictor.pause = True

        # 更新UI状态
        self.run_button.setChecked(False)
        self.save_res_button.setEnabled(True)

        self.pre_video.clear()
        self.res_video.clear()
        self.progress_bar.setValue(0)
        self.Class_num.setText('--')

        self.show_status('Detection terminated')



    def mousePressEvent(self, event):
        p = event.globalPosition()
        globalPos = p.toPoint()
        self.dragPos = globalPos

    def resizeEvent(self, event):
        UIFuncitons.resize_grips(self)

    def closeEvent(self, event):
        # 关闭时保存配置
        sys.exit(0)

if __name__ == "__main__":
    app = QApplication(sys.argv)
    Home = MainWindow()
    Home.show()
    sys.exit(app.exec())